from mbsh.plugins.endoscopic.video_report import TopReport, SkSmallPartActivator, RegexActivator, BrainActivator, \
    LinearFixer

__author__ = 'gpu'

from mbsh.core.models import LabelResult, SmallModel

import unittest
import os

from mbsh import create_app, g_temp_dir, basedir


class VideoReportTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def test_add_lesion(self):
        report = TopReport("病灶导出", 1, 1, 'test_case')
        bbox = [56, 36, 168, 148]
        img_path = basedir + r"\tests\res\video\01.0000000052427.0090.1521592248\imgs\01.0000000047298.0045.1516927596_f0.jpg"
        lesion = LabelResult('', class_index=1, confidence=99, overlay_box=bbox, img_path=img_path)
        clear_confidence = 0.9
        report.add_label(lesion, clear_confidence, part_class_index=1)

        report.add_label(lesion, clear_confidence, part_class_index=0)

        top_list = report.export_top()[0]
        self.assertEqual(2, len(top_list))

        report.add_label(lesion, clear_confidence, part_class_index=0)

        top_list = report.export_top()[0]
        self.assertEqual(2, len(top_list))

        top_list = report.export_top(out_path=os.path.join(g_temp_dir, "test"))[0]
        self.assertEqual(2, len(top_list))
        self.assertEqual(r"d:\znyx\temp\test\1-_c99_f0.jpg", top_list[0].report_path)

    def test_add_parts(self):
        report = TopReport("部位导出", 10, 10, 'test_case')

        part_label0 = LabelResult('', class_index=0, confidence=99)
        clear_confidence = 0.9
        report.add_label(part_label0, clear_confidence)
        top_list = report.export_top()
        self.assertEqual(1, len(top_list[0]))

        part_label2 = LabelResult('', class_index=2, confidence=99)
        report.add_label(part_label2, clear_confidence)

        top_list = report.export_top()
        self.assertEqual(1, len(top_list[0]))
        self.assertEqual(1, len(top_list[2]))

        part_label = LabelResult('', class_index=0, confidence=99)
        report.add_label(part_label, clear_confidence)
        top_list = report.export_top()

        self.assertEqual(2, len(top_list[0]))
        self.assertEqual(1, len(top_list[2]))


class BrainActivatorTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def test_add_case_type(self):
        sm = SmallModel.query.get("胃部分类27")
        sm.check_parents()
        activator = BrainActivator(sm)
        unknow_0 = LabelResult('', class_name="Unknown", class_index=0, confidence=95)
        stomch_1 = LabelResult('', class_name="食管", class_index=1, confidence=80)
        stomch_2 = LabelResult('', class_name="贲门", class_index=2, confidence=80)

        activator.add_label(stomch_1)
        activator.add_label(unknow_0)

        for i in range(5):
            activator.add_label(stomch_1)
        self.assertTrue(stomch_1.has_activate)

        for i in range(4):
            activator.add_label(stomch_2)
        self.assertTrue(stomch_1.has_activate)
        self.assertTrue(stomch_2.has_activate)


class SkLearnActivatorTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def test_add_case_type(self):
        model_name = '胃部分类27'
        sm = SmallModel.query.get(model_name)
        sm.load_model(1, True)
        activator = SkSmallPartActivator(sm, activate_count=3, total_count=5, part_count_dict={8: 2})
        predictions = [0.1] * 28
        predictions[1] = 0.5
        part_label0 = LabelResult(model_name, class_index=1, confidence=99, predictions=predictions)
        activator.add_label(part_label0)
        self.assertFalse(part_label0.has_activate)
        for i in range(13):
            activator.add_label(part_label0)
        self.assertTrue(part_label0.has_activate)

        predictions = [0.1] * 28
        predictions[5] = 0.5
        part_label5 = LabelResult(model_name, class_index=5, confidence=99, predictions=predictions)
        activator.add_label(part_label5)
        self.assertFalse(part_label5.has_activate)

        for i in range(13):
            activator.add_label(part_label5)

        self.assertFalse(part_label5.has_activate)

        predictions = [0.1] * 28
        predictions[8] = 0.9
        part_label8 = LabelResult(model_name, class_index=8, confidence=99, predictions=predictions)
        activator.add_label(part_label8)
        self.assertFalse(part_label8.has_activate)

        for i in range(2):
            activator.add_label(part_label8)

        self.assertTrue(part_label8.has_activate)


class RegexActivatorTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def test_add_case_type(self):
        model_name = '胃部分类27'
        sm = SmallModel.query.get(model_name)
        sm.load_model(1, True)
        activator = RegexActivator(sm, label_regex='(食管,)+(.+,){0,5}贲门', label_index=1)

        part_label0 = LabelResult(model_name, class_index=1, class_name='食管', confidence=99)
        activator.add_label(part_label0)
        self.assertFalse(part_label0.has_activate)
        for i in range(5):
            activator.add_label(part_label0)

        self.assertFalse(part_label0.has_activate)

        label_3 = LabelResult(model_name, class_index=3, class_name='贲门', confidence=99)
        activator.add_label(label_3)
        # 贲门 未激活
        self.assertFalse(label_3.has_activate)

        # 食管 激活
        self.assertTrue(part_label0.has_activate)
        self.assertTrue(label_3.trigger_activate_label.class_name == '食管')

    def test_re(self):
        import re

        self.assertTrue(re.search('(食管,)+(.+,){0,5}贲门', '食管,食管,贲门'))
        self.assertTrue(re.search('(口咽部,)+(.+,){0,5}食管|(食管,)+(.+,){0,5}贲门', '口咽部,无法判断,无法判断,无法判断,食管'))

        self.assertTrue(
            re.match('(胃窦#.+,)+(.+,){0,5}(十二指肠球部,){1,4}|(十二指肠球部,)+(.+,){0,5}十二指肠降部', '十二指肠球部,无法判断,无法判断,无法判断,十二指肠降部'))
        self.assertFalse(
            re.match('^(十二指肠球部,)+(.+,){0,5}十二指肠降部$',
                     '十二指肠球部,无法判断,无法判断,无法判断,十二指肠降部,无法判断'))
        self.assertTrue(
            re.match('(胃窦#.+,)+(.+,){0,5}(十二指肠球部,){1,4}|(十二指肠球部,)+(.+,){0,5}十二指肠降部',
                     '胃窦#小弯,无法判断,胃窦#小弯,无法判断,十二指肠球部,十二指肠球部'))


class LinearFixerTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def test_add_frame(self):
        fixer = LinearFixer(0.5)
        self.assertEqual((0, 40), fixer.add_frame(80))
        self.assertEqual((1, 60), fixer.add_frame(80))
